""" 
    Unit test 
 """


def test_instance_DamageDetectionModel() -> bool:
    """
    尝试实例化模型
    """
    from DamageDetectionModel import DamageDetectionModel

    model = DamageDetectionModel(num_damage_types=5)  # 假设有5种损伤类型
    if model != None:
        return True
    return False



def test_damage_detection_model():
    """ 
        测试模型初始化并构造随机数据测试正向传播
    """
    import torch
    from DamageDetectionModel import DamageDetectionModel
    # 设置随机种子，确保结果可复现
    torch.manual_seed(42)
    
    # 定义输入图像的参数
    batch_size = 2  # 批次大小
    channels = 3    # 彩色图像，3个通道(RGB)
    height = 640    # 图像高度，YOLO通常使用640x640
    width = 640     # 图像宽度
    
    # 创建Mock Tensor作为输入
    # 模拟参考图像 (image_1) 和待配准图像 (image_2)
    image_1 = torch.randn(batch_size, channels, height, width)
    image_2 = torch.randn(batch_size, channels, height, width)
    
    # 打印输入张量信息
    print(f"输入图像1形状: {image_1.shape}")
    print(f"输入图像2形状: {image_2.shape}")
    print(f"输入数据类型: {image_1.dtype}")
    
    # 初始化模型
    num_damage_types = 5
    model = DamageDetectionModel(num_damage_types=num_damage_types)
    model.eval()  # 设置为评估模式
    
    # 进行前向传播
    with torch.no_grad():  # 禁用梯度计算，加快测试速度
        image_merged, outputs = model(image_1, image_2)
    
    # 验证输出
    print("\n测试结果验证:")
    
    # 验证合并图像的形状
    expected_merged_channels = channels * 2  # 因为合并了两个3通道图像
    assert image_merged.shape == (batch_size, expected_merged_channels, height, width), \
        f"合并图像形状不正确，预期 {(batch_size, expected_merged_channels, height, width)}, 实际 {image_merged.shape}"
    print("合并图像形状验证通过")
    
    # 验证YOLO输出
    assert len(outputs) == batch_size, \
        f"YOLO输出数量不正确，预期 {batch_size}, 实际 {len(outputs)}"
    print("YOLO输出数量验证通过")
    
    # 打印输出信息
    print(f"\n合并图像形状: {image_merged.shape}")
    print(f"YOLO输出类型: {type(outputs)}")
    print(f"第一个样本的检测结果: {outputs[0].boxes}")  # 打印第一个样本的检测框信息

def test_all():
    assert test_instance_DamageDetectionModel() == True, "损伤检测模型初始化失败"
    test_damage_detection_model()
